//! linfa机器学习的各种的库的使用方法和封装，为了以后便于使用。
//!
use ndarray::prelude::*;
use rand::prelude::*;

pub mod knn;
mod utils;
pub use utils::*;

/// 把一组数据按照百分比分割成训练数据集和测试数据集.
/// data: 要分隔的数据
/// percent: 是一个（0..1）的小数，不包括0和1， 表示想要把多少分割成训练数据集,
///     剩下的就是测试数据集.
pub fn split_data_to_train_and_test(
    data: &Array2<f64>,
    percent: f64,
) -> anyhow::Result<TrainAndTestData> {
    debug_assert!(percent > 0.0 && percent < 1.0);
    // 计算样本有多少个
    let data_len = data.axis_iter(Axis(0)).len();
    // 需要生成多少个训练样本
    let train_count = (data_len as f64 * percent) as usize;
    // 用来保存生成的索引
    let mut index_vec = Vec::with_capacity(train_count);
    for _ in 0..train_count {
        let mut idx = thread_rng().gen_range(0..data_len);
        while index_vec.contains(&idx) {
            idx = thread_rng().gen_range(0..data_len);
        }
        index_vec.push(idx);
    }
    index_vec.sort();
    // 取出训练样本
    let mut x_train = Array::zeros((0, 3));
    let mut y_train = Array::zeros(0);
    for idx in &index_vec {
        let c = data.index_axis(Axis(0), *idx);
        // 最后一个是target
        x_train.push(Axis(0), c.slice(s![0..(c.len() - 1)]))?;
        y_train.push(Axis(0), c.slice(s![c.len() - 1]))?;
    }

    // 取出测试样本
    let mut x_test = Array::zeros((0, 3));
    let mut y_test = Array::zeros(0);
    for (idx, v) in data.axis_iter(Axis(0)).enumerate() {
        if index_vec.contains(&idx) {
            continue;
        }
        x_test.push(Axis(0), v.slice(s![0..(v.len() - 1)]))?;
        y_test.push(Axis(0), v.slice(s![v.len() - 1]))?;
    }

    Ok(TrainAndTestData {
        x_train,
        y_train,
        x_test,
        y_test,
    })
}

/// 把一组数据划分训练数据和测试数据，用下面的结构体表示
/// 划分的结果。
#[derive(Debug)]
pub struct TrainAndTestData {
    x_train: Array2<f64>,
    y_train: Array1<f64>,
    x_test: Array2<f64>,
    y_test: Array1<f64>,
}

impl TrainAndTestData {
    pub fn x_train(&self) -> &Array2<f64> {
        &self.x_train
    }

    pub fn y_train(&self) -> &Array1<f64> {
        &self.y_train
    }

    pub fn x_test(&self) -> &Array2<f64> {
        &self.x_test
    }

    pub fn y_test(&self) -> &Array1<f64> {
        &self.y_test
    }
}
